import os

import numpy as np
import scipy
import torchvision
from pyts.approximation import PiecewiseAggregateApproximation
from pyts.image import GramianAngularField
from pyts.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
from tqdm import tqdm
from GAF_config import config


def imgshow(data_loader):
    for i, data in enumerate(data_loader):
        images, labels = data
        # 打印数据集中的图片
        img = torchvision.utils.make_grid(images).numpy()
        plt.imshow(np.transpose(img, (1, 2, 0)))
        plt.show()
        break


def pre_data(place, dot, num_train, num_test, id):
    # 读取本地数据
    label = scipy.io.loadmat("./data/" + place + ".mat")
    if not os.path.isdir(config.DATASET + "train\\" + place):
        os.mkdir(config.DATASET + "train\\" + place)
    if not os.path.isdir(config.DATASET + "test\\" + place):
        os.mkdir(config.DATASET + "test\\" + place)
    for i in tqdm(range(num_train)):
        generate_image(place, dot, i, label, True, id)
    for j in tqdm(range(num_test)):
        generate_image(place, dot, j + num_train, label, False, id)
    print("{}-->成功生成{}组训练集数据,{}组测试集数据".format(place, num_train, num_test))


def generate_image(place, dot, i, label, isTrain, id):
    global gram
    x1, y1 = np.linspace(0, dot, dot), \
        label["data"][(i + 1) * dot: (i + 2) * dot, 5]

    if id == 1:
        # PAA
        transformer = PiecewiseAggregateApproximation(window_size=2)
        result = transformer.transform(np.stack((x1, y1), 0).tolist())
        # Scaling in interval [0,1]
        scaler = MinMaxScaler()
        scaled_X = scaler.transform(result)
        arccos_X = np.arccos(scaled_X[1, :])
        field = [a + b for a in arccos_X for b in arccos_X]
        gram = np.cos(field).reshape(-1, 700)

    if id == 2:
        sin_data = y1.reshape(1, -1)
        image_size = 256
        # gasf = GramianAngularField(image_size=image_size, method='summation')
        # sin_gasf = gasf.fit_transform(sin_data)
        gadf = GramianAngularField(image_size=image_size, method='difference')
        sin_gadf = gadf.fit_transform(sin_data)
        gram = sin_gadf[0]

    plt.imshow(gram)
    plt.axis('off')
    if isTrain:
        plt.savefig('./data/train/{}/{}_{}.png'.format(place, place, (i + 1)), bbox_inches="tight", pad_inches=0.0)
    else:
        plt.savefig('./data/test/{}/{}_{}.png'.format(place, place, (i + 1)), bbox_inches="tight", pad_inches=0.0)
